baml-language/fix: use build_plan for llm functions, delete execute#3279
baml-language/fix: use build_plan for llm functions, delete execute#3279
Conversation
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
📝 WalkthroughWalkthroughThe changes refactor the LLM execution flow from direct client execution to a step-based orchestration model. The Changes
Sequence Diagram(s)sequenceDiagram
participant Caller
participant CallLLMFunc as call_llm_function
participant Client
participant Orchestration as Orchestration Step
participant ErrorHandler
Caller->>CallLLMFunc: call_llm_function<T>(context)
CallLLMFunc->>Client: build_plan()
Client-->>CallLLMFunc: steps: [OrchestrationStep]
CallLLMFunc->>Client: advance_round_robin()
Client-->>CallLLMFunc: state updated
loop For each step in steps
CallLLMFunc->>CallLLMFunc: check step.delay_ms
alt delay_ms > 0
CallLLMFunc->>CallLLMFunc: sleep(delay_ms)
end
CallLLMFunc->>Orchestration: execute_step(step, context)
alt Step succeeds
Orchestration-->>CallLLMFunc: result: T
CallLLMFunc-->>Caller: return result
else Step fails
Orchestration-->>ErrorHandler: exception caught
ErrorHandler-->>CallLLMFunc: continue to next step
end
end
alt All steps failed
CallLLMFunc-->>CallLLMFunc: throw DevOther("All orchestration steps failed")
CallLLMFunc-->>Caller: error
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes 🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 11ce9f11b9
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| let result: T = client.execute(context, 0); | ||
| result | ||
| let steps = client.build_plan(); | ||
| client.advance_round_robin(); |
There was a problem hiding this comment.
Advance round-robin counters only for attempted steps
call_llm_function now mutates round-robin state before any step runs, which breaks fallback behavior when an earlier branch succeeds. Because build_plan() flattens all fallback branches but the loop returns on the first successful step, many planned steps are never attempted; however client.advance_round_robin() still increments counters in those untouched subtrees, so future calls can skip providers that were never actually used. This is a regression from the previous execute-on-visit behavior and changes routing deterministically for Fallback[..., RoundRobin[...]] clients.
Useful? React with 👍 / 👎.
11ce9f1 to
fbe1d47
Compare
Binary size checks passed✅ 7 passed
Generated by |
fbe1d47 to
b11e050
Compare
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
baml_language/crates/bex_engine/tests/orchestration.rs (1)
101-129: 🧹 Nitpick | 🔵 TrivialRemove the ignored debug probe or make it assert on
FB.
diag_fallback_fieldsnever runs, has no assertions, andcheck()returns[A, B]instead of exercisingFB, so it does not inspect fallback fields at all. This reads like leftover debugging code rather than regression coverage.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 7fcc270b-cb1e-4061-b5f4-edcb5669d19a
📒 Files selected for processing (4)
baml_language/crates/baml_builtins2/baml_std/baml/llm.bamlbaml_language/crates/baml_builtins2/baml_std/baml/llm_types.bamlbaml_language/crates/baml_compiler2_tir/src/builder.rsbaml_language/crates/bex_engine/tests/orchestration.rs
| function advance_round_robin(self) -> void { | ||
| match (self.client_type) { | ||
| ClientType.Primitive => {}, | ||
| ClientType.Fallback => { | ||
| for (let sub in self.sub_clients) { | ||
| sub.advance_round_robin(); | ||
| } | ||
| }, | ||
| ClientType.RoundRobin => { | ||
| self.counter += 1; | ||
| }, | ||
| } | ||
| } |
There was a problem hiding this comment.
Advance the selected round-robin child too.
build_plan_with_state() consumes the chosen child subtree after computing idx, so nested round-robin nodes under that child can participate in the plan. advance_round_robin() only increments self.counter, though, so those nested counters never advance and the same inner branch is reused whenever this outer RR picks that child again. Mirror the same child selection here and recurse into the chosen sub-client as well.
| let result: T = execute_step(step, context) catch (e) { | ||
| _ => { continue; } | ||
| }; | ||
| return result; | ||
| } | ||
|
|
||
| throw root.errors.DevOther { message: "All orchestration steps failed" }; |
There was a problem hiding this comment.
Don't treat every execute_step() failure as retryable.
execute_step() can fail with deterministic configuration errors too, not just provider failures. Catching _ here retries the next step and eventually replaces the real cause with "All orchestration steps failed", which can also send extra requests for a non-retryable bug. Restrict the fallback path to step-local provider/HTTP failures and rethrow the rest.
| Definition::Let(let_loc) => { | ||
| // Determine type from the let-binding's origin. | ||
| let db = self.context.db(); | ||
| let item_tree = | ||
| baml_compiler2_hir::file_item_tree(db, let_loc.file(db)); | ||
| let let_data = &item_tree[let_loc.id(db)]; | ||
| match let_data.origin { | ||
| baml_compiler2_ast::ast::LetOrigin::Client => { | ||
| // client<llm> declarations produce Client instances. | ||
| Ty::Class(crate::ty::QualifiedTypeName::new( | ||
| baml_base::Name::new("baml"), | ||
| vec![baml_base::Name::new("llm")], | ||
| baml_base::Name::new("Client"), | ||
| )) | ||
| } | ||
| baml_compiler2_ast::ast::LetOrigin::RetryPolicy => { | ||
| // retry_policy declarations produce RetryPolicy instances. | ||
| Ty::Class(crate::ty::QualifiedTypeName::new( | ||
| baml_base::Name::new("baml"), | ||
| vec![baml_base::Name::new("llm")], | ||
| baml_base::Name::new("RetryPolicy"), | ||
| )) | ||
| } | ||
| _ => Ty::Unknown, | ||
| } |
There was a problem hiding this comment.
Handle let-bound globals in package-qualified paths too.
This branch only fixes bare identifiers through infer_single_name(). root.A.build_plan() and package-qualified client/retry-policy references still go through resolve_package_item(), which only recognizes Definition::Function, so they continue to infer as Ty::Unknown and member lookup fails. Please share the same Definition::Let → type mapping with the multi-segment/package path resolver.
Summary by CodeRabbit
Refactor
Tests